{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent Based Models\n",
    "\n",
    "Code examples from [Think Complexity, 2nd edition](http://greenteapress.com/wp/complexity2), Chapter 9\n",
    "\n",
    "Copyright 2016 Allen Downey, [MIT License](http://opensource.org/licenses/MIT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import thinkplot\n",
    "from thinkstats2 import Cdf\n",
    "from thinkstats2 import RandomSeed\n",
    "\n",
    "from matplotlib import rc\n",
    "rc('animation', html='html5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Schelling's model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Cell2D import Cell2D, Cell2DViewer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`locs_where` is a wrapper on `np.where` that returns results in an array (rather than two tuples)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def locs_where(condition):\n",
    "    \"\"\"Find cells where a logical array is True.\n",
    "    \n",
    "    condition: logical array\n",
    "    \n",
    "    returns: array with one set of coordinates per row\n",
    "    \"\"\"\n",
    "    return np.transpose(np.nonzero(condition))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's my implementation of Schelling's model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import correlate2d\n",
    "\n",
    "class Schelling(Cell2D):\n",
    "    \"\"\"Represents a grid of Schelling agents.\"\"\"\n",
    "    \n",
    "    options = dict(mode='same', boundary='wrap')\n",
    "\n",
    "    kernel = np.array([[1, 1, 1],\n",
    "                       [1, 0, 1],\n",
    "                       [1, 1, 1]], dtype=np.int8)\n",
    "    \n",
    "    def __init__(self, n, m=None, p=0.5):\n",
    "        \"\"\"Initializes the attributes.\n",
    "\n",
    "        n: number of rows\n",
    "        m: number of columns\n",
    "        p: threshold on the fraction of similar neighbors\n",
    "        \"\"\"\n",
    "        self.p = p\n",
    "        m = n if m is None else m\n",
    "        # 0 is empty, 1 is red, 2 is blue\n",
    "        choices = [0, 1, 2]\n",
    "        probs = [0.1, 0.45, 0.45]\n",
    "        self.array = np.random.choice(choices, (n, m), p=probs).astype(np.int8)\n",
    "\n",
    "    def count_neighbors(self):\n",
    "        \"\"\"Surveys neighboring cells.\n",
    "        \n",
    "        returns: tuple of (occupied, frac_red, frac_same)\n",
    "        \n",
    "        occupied: logical array indicating occupied cells\n",
    "        frac_red: fraction of red neighbors around each cell\n",
    "        frac_same: for each cell, the fraction of similar neighbors\n",
    "        \"\"\"\n",
    "        a = self.array\n",
    "        \n",
    "        red = a==1\n",
    "        blue = a==2\n",
    "        occupied = a!=0\n",
    "\n",
    "        # count red neighbors and all neighbors\n",
    "        num_red = correlate2d(red, self.kernel, **self.options)\n",
    "        num_blue = correlate2d(blue, self.kernel, **self.options)\n",
    "        num_neighbors = num_red + num_blue\n",
    "\n",
    "        # compute fraction of similar neighbors\n",
    "        frac_red = num_red / num_neighbors\n",
    "        frac_blue = num_blue / num_neighbors\n",
    "        frac_same = np.where(red, frac_red, frac_blue)\n",
    "        \n",
    "        # no neighbors is considered the same as no similar neighbors \n",
    "        frac_same[num_neighbors == 0] = 0\n",
    "        frac_red[num_neighbors == 0] = 0\n",
    "        \n",
    "        return occupied, frac_red, frac_same\n",
    "\n",
    "    def segregation(self):\n",
    "        \"\"\"Computes the average fraction of similar neighbors.\n",
    "        \n",
    "        returns: fraction of similar neighbors, averaged over cells\n",
    "        \"\"\"\n",
    "        occupied, _, frac_same = self.count_neighbors()\n",
    "        return np.sum(frac_same) / np.sum(occupied)\n",
    "    \n",
    "    def step(self):\n",
    "        \"\"\"Executes one time step.\n",
    "                \n",
    "        returns: fraction of similar neighbors, averaged over cells\n",
    "        \"\"\"\n",
    "        a = self.array\n",
    "        \n",
    "        # find the unhappy cells\n",
    "        occupied, _, frac_same = self.count_neighbors()\n",
    "        unhappy_locs = locs_where(occupied & (frac_same < self.p))\n",
    "\n",
    "        # find the empty cells\n",
    "        empty = a==0\n",
    "        num_empty = np.sum(empty)\n",
    "        empty_locs = locs_where(empty)\n",
    "\n",
    "        # shuffle the unhappy cells\n",
    "        if len(unhappy_locs):\n",
    "            np.random.shuffle(unhappy_locs)\n",
    "            \n",
    "        # for each unhappy cell, choose a random destination\n",
    "        for source in unhappy_locs:\n",
    "            i = np.random.randint(len(empty_locs))\n",
    "            dest = tuple(empty_locs[i])\n",
    "\n",
    "            # move\n",
    "            a[dest] = a[tuple(source)]\n",
    "            a[tuple(source)] = 0\n",
    "            empty_locs[i] = source\n",
    "        \n",
    "        num_empty2 = np.sum(a==0)\n",
    "        assert num_empty == num_empty2\n",
    "        \n",
    "        # return the average fraction of similar neighbors\n",
    "        return np.sum(frac_same) / np.sum(occupied)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_cmap` makes a custom color map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "\n",
    "def make_cmap(color_dict, vmax=None, name='mycmap'):\n",
    "    \"\"\"Makes a custom color map.\n",
    "    \n",
    "    color_dict: map from numbers to colors\n",
    "    vmax: high end of the range, \n",
    "    name: string name for map\n",
    "    \n",
    "    If vmax is None, uses the max value from color_dict\n",
    "    \n",
    "    returns: pyplot color map\n",
    "    \"\"\"\n",
    "    if vmax is None:\n",
    "        vmax = max(color_dict.keys())\n",
    "    \n",
    "    colors = [(value/vmax, color) for value, color in color_dict.items()]\n",
    "    \n",
    "    cmap = LinearSegmentedColormap.from_list(name, colors)\n",
    "    \n",
    "    return cmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `SchellingViewer` is a basic `Cell2DViewer`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SchellingViewer(Cell2DViewer):\n",
    "    # colors from http://colorbrewer2.org/#type=qualitative&scheme=Accent&n=5\n",
    "    colors = ['#7fc97f','#beaed4','#fdc086','#ffff99','#386cb0']\n",
    "    cmap = make_cmap({0:'white', 1:colors[2], 2:colors[4]})\n",
    "    options = dict(interpolation='none', alpha=0.8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a small example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = Schelling(n=10)\n",
    "viewer = SchellingViewer(grid)\n",
    "viewer.draw()\n",
    "grid.segregation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's an animation for a bigger example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = Schelling(n=100)\n",
    "viewer = SchellingViewer(grid)\n",
    "anim = viewer.animate(frames=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The degree of segregation increases quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following figure shows the process after 2 and 10 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(17)\n",
    "\n",
    "grid = Schelling(n=100, p=0.3)\n",
    "print(grid.segregation())\n",
    "\n",
    "thinkplot.preplot(cols=3)\n",
    "viewer = SchellingViewer(grid)\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(2)\n",
    "for i in range(2):\n",
    "    viewer.step()\n",
    "print(grid.segregation())\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(3)\n",
    "for i in range(8):\n",
    "    viewer.step()\n",
    "print(grid.segregation())\n",
    "viewer.draw()\n",
    "\n",
    "plt.savefig('chap09-1.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's how segregation in steady state relates to `p`, the threshold on the fraction of similar neighbors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(17)\n",
    "\n",
    "for p in [0.5, 0.4, 0.3, 0.2]:\n",
    "    grid = Schelling(n=100, p=p)\n",
    "    segs = [grid.step() for i in range(20)]\n",
    "    thinkplot.plot(segs, label='p = %.1f' % p)\n",
    "    print(p, segs[-1], segs[-1] - p)\n",
    "    \n",
    "thinkplot.config(xlabel='Time steps', ylabel='Segregation',\n",
    "                loc='lower right', ylim=[0, 1])\n",
    "\n",
    "plt.savefig('chap09-2.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At `p=0.3`, there is a striking difference between the level that would make people happy, at only 30%, and the level they actually get, around 80%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Experiment with different starting conditions: for example, more or fewer empty cells, or unequal numbers of red and blue agents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sugarscape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_locs` takes the dimensions of the grid and returns an array where each row is a coordinate in the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_locs(n, m):\n",
    "    \"\"\"Makes array where each row is an index in an `n` by `m` grid.\n",
    "    \n",
    "    n: int number of rows\n",
    "    m: int number of cols\n",
    "    \n",
    "    returns: NumPy array\n",
    "    \"\"\"\n",
    "    left = np.repeat(np.arange(m), n)\n",
    "    right = np.tile(np.arange(n), m)\n",
    "    return np.transpose([left, right])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_locs(2, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_visible_locs` takes the range of an agents vision and returns an array where each row is the coordinate of a visible cell.\n",
    "\n",
    "The cells are at increasing distances.  The cells at each distance are shuffled. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_visible_locs(vision):\n",
    "    \"\"\"Computes the kernel of visible cells.\n",
    "        \n",
    "    vision: int distance\n",
    "    \"\"\"\n",
    "    def make_array(d):\n",
    "        \"\"\"Generates visible cells with increasing distance.\"\"\"\n",
    "        a = np.array([[-d, 0], [d, 0], [0, -d], [0, d]])\n",
    "        np.random.shuffle(a)\n",
    "        return a\n",
    "                     \n",
    "    arrays = [make_array(d) for d in range(1, vision+1)]\n",
    "    return np.vstack(arrays)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_visible_locs(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's my implementation of Sugarscape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sugarscape(Cell2D):\n",
    "    \"\"\"Represents an Epstein-Axtell Sugarscape.\"\"\"\n",
    "    \n",
    "    def __init__(self, n, **params):\n",
    "        \"\"\"Initializes the attributes.\n",
    "\n",
    "        n: number of rows and columns\n",
    "        params: dictionary of parameters\n",
    "        \"\"\"\n",
    "        self.n = n\n",
    "        self.params = params\n",
    "        \n",
    "        # track variables\n",
    "        self.agent_count_seq = []\n",
    "    \n",
    "        # make the capacity array\n",
    "        self.capacity = self.make_capacity()\n",
    "        \n",
    "        # initially all cells are at capacity\n",
    "        self.array = self.capacity.copy()\n",
    "        \n",
    "        # make the agents\n",
    "        self.make_agents()\n",
    "        \n",
    "    def make_capacity(self):\n",
    "        \"\"\"Makes the capacity array.\"\"\"\n",
    "        \n",
    "        # compute the distance of each cell from the peaks. \n",
    "        x = np.arange(self.n)\n",
    "        rows, cols = np.meshgrid(x, x, indexing='ij')\n",
    "        \n",
    "        # each cell in `rows` contains its own `i` coordinate\n",
    "        # each cell in `cols` contains its `j` coordinate\n",
    "        dist1 = np.hypot(rows-15, cols-15)\n",
    "        dist2 = np.hypot(rows-35, cols-35)\n",
    "        \n",
    "        # each cell in `dist` contains its distance to the closer peak\n",
    "        dist = np.minimum(dist1, dist2)\n",
    "        \n",
    "        # cells in the capacity array are set according to dist from peak\n",
    "        a = np.zeros((self.n, self.n), np.float)\n",
    "        a[dist<21] = 1\n",
    "        a[dist<16] = 2\n",
    "        a[dist<11] = 3\n",
    "        a[dist<6] = 4\n",
    "\n",
    "        return a\n",
    "        \n",
    "    def make_agents(self):\n",
    "        \"\"\"Makes the agents.\"\"\"\n",
    "        \n",
    "        # determine where the agents start and generate locations\n",
    "        n, m = self.params.get('starting_box', self.array.shape)\n",
    "        locs = make_locs(n, m)\n",
    "        np.random.shuffle(locs)\n",
    "\n",
    "        # make the agents\n",
    "        num_agents = self.params.get('num_agents', 400)\n",
    "        self.agents = [Agent(locs[i], self.params) \n",
    "                       for i in range(num_agents)]\n",
    "        \n",
    "        # keep track of which cells are occupied\n",
    "        self.occupied = set(agent.loc for agent in self.agents)\n",
    "            \n",
    "    def grow(self):\n",
    "        \"\"\"Adds sugar to all cells and caps them by capacity.\"\"\"\n",
    "        grow_rate = self.params.get('grow_rate', 1)\n",
    "        self.array = np.minimum(self.array + grow_rate, self.capacity)\n",
    "        \n",
    "    def look_around(self, center, vision):\n",
    "        \"\"\"Finds the visible cell with the most sugar.\n",
    "        \n",
    "        center: tuple, coordinates of the center cell\n",
    "        vision: int, maximum visible distance\n",
    "        \n",
    "        returns: tuple, coordinates of best cell\n",
    "        \"\"\"\n",
    "        # find all visible cells\n",
    "        locs = make_visible_locs(vision)\n",
    "        locs = (locs + center) % self.n\n",
    "        \n",
    "        # convert rows of the array to tuples\n",
    "        locs = [tuple(loc) for loc in locs]\n",
    "        \n",
    "        # select unoccupied cells\n",
    "        empty_locs = [loc for loc in locs if loc not in self.occupied]\n",
    "        \n",
    "        # if all visible cells are occupied, stay put\n",
    "        if len(empty_locs) == 0:\n",
    "            return center\n",
    "        \n",
    "        # look up the sugar level in each cell\n",
    "        t = [self.array[loc] for loc in empty_locs]\n",
    "        \n",
    "        # find the best one and return it\n",
    "        # (in case of tie, argmax returns the first, which\n",
    "        # is the closest)\n",
    "        i = np.argmax(t)\n",
    "        return empty_locs[i]\n",
    "    \n",
    "    def harvest(self, loc):\n",
    "        \"\"\"Removes and returns the sugar from `loc`.\n",
    "        \n",
    "        loc: tuple coordinates\n",
    "        \"\"\"\n",
    "        sugar = self.array[loc]\n",
    "        self.array[loc] = 0\n",
    "        return sugar\n",
    "    \n",
    "    def step(self):\n",
    "        \"\"\"Executes one time step.\"\"\"\n",
    "        replace = self.params.get('replace', False)\n",
    "        \n",
    "        # loop through the agents in random order\n",
    "        random_order = np.random.permutation(self.agents)\n",
    "        for agent in random_order:\n",
    "            \n",
    "            # mark the current cell unoccupied\n",
    "            self.occupied.remove(agent.loc)\n",
    "            \n",
    "            # execute one step\n",
    "            agent.step(self)\n",
    "\n",
    "            # if the agent is dead, remove from the list\n",
    "            if agent.is_starving() or agent.is_old():\n",
    "                self.agents.remove(agent)\n",
    "                if replace:\n",
    "                    self.add_agent()\n",
    "            else:\n",
    "                # otherwise mark its cell occupied\n",
    "                self.occupied.add(agent.loc)\n",
    "\n",
    "        # update the time series\n",
    "        self.agent_count_seq.append(len(self.agents))\n",
    "        \n",
    "        # grow back some sugar\n",
    "        self.grow()\n",
    "        return len(self.agents)\n",
    "    \n",
    "    def add_agent(self):\n",
    "        \"\"\"Generates a new random agent.\n",
    "                \n",
    "        returns: new Agent\n",
    "        \"\"\"\n",
    "        new_agent = Agent(self.random_loc(), self.params)\n",
    "        self.agents.append(new_agent)\n",
    "        self.occupied.add(new_agent.loc)\n",
    "        return new_agent\n",
    "    \n",
    "    def random_loc(self):\n",
    "        \"\"\"Choose a random unoccupied cell.\n",
    "        \n",
    "        returns: tuple coordinates\n",
    "        \"\"\"\n",
    "        while True:\n",
    "            loc = tuple(np.random.randint(self.n, size=2))\n",
    "            if loc not in self.occupied:\n",
    "                return loc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's a viewer that displays the state of Sugarscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SugarscapeViewer(Cell2DViewer):\n",
    "    \"\"\"Generates visualization and animation of Sugarscape.\"\"\"\n",
    "    \n",
    "    cmap = plt.get_cmap('YlOrRd')\n",
    "    \n",
    "    options = dict(interpolation='none', alpha=0.8,\n",
    "                   vmin=0, vmax=9)\n",
    "    \n",
    "    def draw(self, grid=False):\n",
    "        \"\"\"Draws the array and any other elements.\n",
    "        \n",
    "        grid: boolean, whether to draw grid lines\n",
    "        \"\"\"\n",
    "        self.draw_array(self.viewee.array, origin='lower')\n",
    "        self.draw_agents()\n",
    "\n",
    "    def draw_agents(self):\n",
    "        \"\"\"Plots the agents.\n",
    "        \"\"\"\n",
    "        xs, ys = self.get_coords()\n",
    "        self.points = plt.plot(xs, ys, '.', color='red')[0]\n",
    "        \n",
    "    def animate_func(self, i):\n",
    "        \"\"\"Draws one frame of the animation.\"\"\"\n",
    "        Cell2DViewer.animate_func(self, i)\n",
    "        xs, ys = self.get_coords()\n",
    "        self.points.set_data(np.array([xs, ys]))\n",
    "        return self.im, self.points\n",
    "    \n",
    "    def get_coords(self):\n",
    "        \"\"\"Gets the coordinates of the agents.\n",
    "        \n",
    "        Transforms from (row, col) to (x, y).\n",
    "        \n",
    "        returns: tuple of sequences, (xs, ys)\n",
    "        \"\"\"\n",
    "        agents = self.viewee.agents\n",
    "        rows, cols = np.transpose([agent.loc for agent in agents])\n",
    "        xs = cols + 0.5\n",
    "        ys = rows + 0.5\n",
    "        return xs, ys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's my implementation of the agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    \n",
    "    def __init__(self, loc, params):\n",
    "        \"\"\"Creates a new agent at the given location.\n",
    "        \n",
    "        loc: tuple coordinates\n",
    "        params: dictionary of parameters\n",
    "        \"\"\"\n",
    "        self.loc = tuple(loc)\n",
    "        self.age = 0\n",
    "\n",
    "        # extract the parameters\n",
    "        max_vision = params.get('max_vision', 6)\n",
    "        max_metabolism = params.get('max_metabolism', 4)\n",
    "        min_lifespan = params.get('min_lifespan', 10000)\n",
    "        max_lifespan = params.get('max_lifespan', 10000)\n",
    "        min_sugar = params.get('min_sugar', 5)\n",
    "        max_sugar = params.get('max_sugar', 25)\n",
    "        \n",
    "        # choose attributes\n",
    "        self.vision = np.random.random_integers(max_vision)\n",
    "        self.metabolism = np.random.uniform(1, max_metabolism)\n",
    "        self.lifespan = np.random.uniform(min_lifespan, max_lifespan)\n",
    "        self.sugar = np.random.uniform(min_sugar, max_sugar)\n",
    "\n",
    "    def step(self, env):\n",
    "        \"\"\"Look around, move, and harvest.\n",
    "        \n",
    "        env: Sugarscape\n",
    "        \"\"\"\n",
    "        self.loc = env.look_around(self.loc, self.vision)\n",
    "        self.sugar += env.harvest(self.loc) - self.metabolism\n",
    "        self.age += 1\n",
    "\n",
    "    def is_starving(self):\n",
    "        \"\"\"Checks if sugar has gone negative.\"\"\"\n",
    "        return self.sugar < 0\n",
    "    \n",
    "    def is_old(self):\n",
    "        \"\"\"Checks if lifespan is exceeded.\"\"\"\n",
    "        return self.age > self.lifespan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example with `n=50`, starting with 400 agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "env = Sugarscape(50, num_agents=400)\n",
    "viewer = SugarscapeViewer(env)\n",
    "viewer.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distribution of vision is uniform from 1 to 6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = Cdf(agent.vision for agent in env.agents)\n",
    "thinkplot.Cdf(cdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distribution of metabolism is uniform from 1 to 4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = Cdf(agent.metabolism for agent in env.agents)\n",
    "thinkplot.Cdf(cdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The distribution of initial endowment of sugar is uniform from 5 to 25."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = Cdf(agent.sugar for agent in env.agents)\n",
    "thinkplot.Cdf(cdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the animation looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = viewer.animate(frames=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of agents levels off at the \"carrying capacity\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(env.agents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thinkplot.plot(env.agent_count_seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This figure shows the state of the system after 2 and 10 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(17)\n",
    "\n",
    "env = Sugarscape(50, num_agents=400)\n",
    "viewer = SugarscapeViewer(env)\n",
    "\n",
    "thinkplot.preplot(cols=3)\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(2)\n",
    "for i in range(2):\n",
    "    viewer.step()\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(3)\n",
    "for i in range(98):\n",
    "    viewer.step()\n",
    "viewer.draw()\n",
    "\n",
    "plt.savefig('chap09-3.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Experiment with different numbers of agents.  Try increasing or decreasing their vision or metabolism, and see what effect is has on carrying capacity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sugarscape with finite lifespans"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we start with 250 agents, with lifetimes from 60 to 100, and replacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = Sugarscape(50, num_agents=250,\n",
    "                 min_lifespan=60, max_lifespan=100, \n",
    "                 replace=True)\n",
    "\n",
    "viewer = SugarscapeViewer(env)\n",
    "anim = viewer.animate(frames=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the animation looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After 100 time steps, the distribution of wealth is skewed to the right.  Most agents have very little sugar, but a few have a lot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf = Cdf(agent.sugar for agent in env.agents)\n",
    "thinkplot.Cdf(cdf)\n",
    "thinkplot.config(xlabel='Wealth', ylabel='CDF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf.Percentiles([25, 50, 75, 90])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting with the same parameters, I'll run the model 500 steps, recording the distribution of wealth after each 100 steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(17)\n",
    "\n",
    "env = Sugarscape(50, num_agents=250,\n",
    "                 min_lifespan=60, max_lifespan=100, \n",
    "                 replace=True)\n",
    "\n",
    "cdfs = []\n",
    "for i in range(5):\n",
    "    [env.step() for i in range(100)]\n",
    "    cdf = Cdf(agent.sugar for agent in env.agents)\n",
    "    cdfs.append(cdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After about 200 steps, the distribution is stationary (doesn't change over time).\n",
    "\n",
    "On a log scale, it is approximately normal, possibly with a truncated right tail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thinkplot.preplot(cols=2)\n",
    "\n",
    "thinkplot.Cdfs(cdfs[:-1], color='gray', alpha=0.3)\n",
    "thinkplot.Cdf(cdfs[-1])\n",
    "thinkplot.config(xlabel='Weath', ylabel='CDF')\n",
    "\n",
    "thinkplot.subplot(2)\n",
    "thinkplot.Cdfs(cdfs[:-1], color='gray', alpha=0.3)\n",
    "thinkplot.Cdf(cdfs[-1])\n",
    "thinkplot.config(xlabel='Weath', ylabel='CDF', xscale='log')\n",
    "\n",
    "plt.savefig('chap09-4.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Experiment with different starting conditions and agents with different vision, metabolism, and lifespan.  What effect do these changes have on the distribution of wealth?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Migration in waves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we start with all agents in the lower left, they propagate up and to the right in waves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RandomSeed(17)\n",
    "\n",
    "env = Sugarscape(50, num_agents=300, starting_box=(20, 20), max_vision=16)\n",
    "    \n",
    "viewer = SugarscapeViewer(env)\n",
    "anim = viewer.animate(frames=20, interval=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like after 6 and 12 steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = Sugarscape(50, num_agents=300, starting_box=(20, 20), max_vision=16)\n",
    "viewer = SugarscapeViewer(env)\n",
    "\n",
    "thinkplot.preplot(cols=3)\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(2)\n",
    "for i in range(6):\n",
    "    viewer.step()\n",
    "viewer.draw()\n",
    "\n",
    "thinkplot.subplot(3)\n",
    "for i in range(6):\n",
    "    viewer.step()\n",
    "viewer.draw()\n",
    "\n",
    "plt.savefig('chap09-5.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example is interesting because the waves move diagonally, unlike the agents, who can only move up or to the right.  They are similar in some ways to gliders and other Game of Life spaceships. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Again, experiment with different starting conditions and see what effect they have on the wave behavior."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Bill Bishop, author of *The Big Sort*, argues that\n",
    "American society is increasingly segregated by political\n",
    "opinion, as people choose to live among like-minded neighbors.\n",
    "\n",
    "The mechanism Bishop hypothesizes is not that people, like the agents\n",
    "in Schelling's model, are more likely to move if they are\n",
    "isolated, but that when they move for any reason, they are\n",
    "likely to choose a neighborhood with people like themselves.\n",
    "\n",
    "Write a version of Schelling's model to simulate\n",
    "this kind of behavior and see if it yields similar degrees of\n",
    "segregation.\n",
    "\n",
    "There are several ways you can model Bishop's hypothesis.  In my\n",
    "implementation, a random selection of agents moves during each step.\n",
    "Each agent considers `k` randomly-chosen empty locations and\n",
    "chooses the one with the highest fraction of similar neighbors.\n",
    "How does the degree of segregation depend on `k`?\n",
    "\n",
    "You should be able to implement this model by inheriting from\n",
    "`Schelling` and overriding `__init__` and `step`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "class BigSort(Schelling):\n",
    "    \n",
    "    def __init__(self, n, m=None, num_comps=2):\n",
    "        \"\"\"Initializes the attributes.\n",
    "\n",
    "        n: number of rows\n",
    "        m: number of columns\n",
    "        p: threshold on the fraction of similar neighbors\n",
    "        \"\"\"\n",
    "        self.num_comps = num_comps\n",
    "        m = n if m is None else m\n",
    "        choices = [0, 1, 2]\n",
    "        probs = [0.1, 0.45, 0.45]\n",
    "        self.array = np.random.choice(choices, (n, m), p=probs).astype(np.int8)\n",
    "\n",
    "    def step(self, p=0.1):\n",
    "        \"\"\"Executes one time step.\n",
    "        \n",
    "        p: fraction of agents who move\n",
    "        \n",
    "        returns: fraction of similar neighbors, averaged over cells\n",
    "        \"\"\"\n",
    "        a = self.array\n",
    "        \n",
    "        # find the empty cells\n",
    "        empty = a==0\n",
    "        num_empty = np.sum(empty)\n",
    "        empty_locs = locs_where(empty)\n",
    "        \n",
    "        # choose the cells that are moving\n",
    "        occupied, frac_red, frac_same = self.count_neighbors()\n",
    "        r = np.random.random(a.shape)\n",
    "        unhappy_locs = locs_where(occupied & (r < p))\n",
    "\n",
    "        # shuffle the unhappy cells\n",
    "        if len(unhappy_locs):\n",
    "            np.random.shuffle(unhappy_locs)\n",
    "            \n",
    "        # for each unhappy cell, choose a destination and move\n",
    "        for loc in unhappy_locs:\n",
    "            source = tuple(loc)\n",
    "            \n",
    "            # make a list of random choices\n",
    "            choices = []\n",
    "            start = np.random.randint(len(empty_locs))\n",
    "            for i in range(start, start+self.num_comps):\n",
    "                i %= len(empty_locs)\n",
    "                dest = tuple(empty_locs[i])\n",
    "                frac = frac_red[dest]\n",
    "                choices.append((frac, i, dest))\n",
    "            \n",
    "            # choose a destination\n",
    "            if a[source] == 1:\n",
    "                # if red, maximize the fraction of red\n",
    "                frac, i, dest = max(choices)\n",
    "            else:\n",
    "                # if blue, minimize\n",
    "                frac, i, dest = min(choices)\n",
    "\n",
    "            # move\n",
    "            a[dest] = a[source]\n",
    "            a[source] = 0\n",
    "            empty_locs[i] = loc\n",
    "        \n",
    "        # check that the number of empty cells hasn't changed\n",
    "        num_empty2 = np.sum(a==0)\n",
    "        assert num_empty == num_empty2\n",
    "        \n",
    "        # return the average fraction of similar neighbors\n",
    "        return np.sum(frac_same) / np.sum(occupied)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Here's a small example.\n",
    "\n",
    "grid = BigSort(n=10)\n",
    "viewer = SchellingViewer(grid)\n",
    "viewer.draw()\n",
    "grid.segregation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a test of the `step` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# And a test of the `step` method\n",
    "\n",
    "grid.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Here's what the animation looks like with `num_comps=2`\n",
    "\n",
    "grid = BigSort(n=100)\n",
    "viewer = SchellingViewer(grid)\n",
    "anim = viewer.animate(frames=60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Looks like things get pretty segregated, although not as\n",
    "# quickly as in the Schelling model.\n",
    "\n",
    "# Here's how the degree of segregation depends on the number\n",
    "# of locations each agent considers.\n",
    "\n",
    "for num_comps in [4, 3, 2]:\n",
    "    grid = BigSort(n=100, num_comps=num_comps)\n",
    "    segs = [grid.step() for i in range(60)]\n",
    "    thinkplot.plot(segs, label=num_comps)\n",
    "    print(num_comps, segs[-1])\n",
    "    \n",
    "thinkplot.config(xlabel='Steps', ylabel='Segregation', loc='upper left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Here's what the final state looks like with `num_comps=4`\n",
    "\n",
    "# Unlike the Schelling model, the Big Sort model doesn't seem \n",
    "# to aggregate the empty cells.\n",
    "\n",
    "viewer = SchellingViewer(grid)\n",
    "viewer.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** In the first version of Sugarscape, we never add agents, so once the population falls, it never recovers.  In the second version, we only replace agents when they die, so the population is constant.  Now let's see what happens if we add some \"population pressure\".\n",
    "\n",
    "Write a version of Sugarscape that adds a new agent at the end of every step.  Add code to compute the average vision and the average metabolism of the agents at the end of each step.  Run the model for a few hundred steps and plot the population over time, as well as the average vision and average metabolism.\n",
    "\n",
    "You should be able to implement this model by inheriting from\n",
    "`Sugarscape` and overriding `__init__` and `step`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "class EvoSugarscape(Sugarscape):\n",
    "    \"\"\"Represents an Epstein-Axtell Sugarscape.\"\"\"\n",
    "    \n",
    "    def __init__(self, n, **params):\n",
    "        \"\"\"Initializes the attributes.\n",
    "\n",
    "        n: number of rows and columns\n",
    "        params: dictionary of parameters\n",
    "        \"\"\"\n",
    "        Sugarscape.__init__(self, n, **params)\n",
    "        \n",
    "        # track variables\n",
    "        self.avg_vision_seq = []\n",
    "        self.avg_metabolism_seq = []\n",
    "    \n",
    "    def step(self):\n",
    "        \"\"\"Executes one time step.\"\"\"\n",
    "        Sugarscape.step(self)\n",
    "        \n",
    "        # average vision\n",
    "        avg_vision = np.mean([agent.vision for agent in self.agents])\n",
    "        self.avg_vision_seq.append(avg_vision)\n",
    "        \n",
    "        # average metabolism\n",
    "        avg_metabolism = np.mean([agent.metabolism for agent in self.agents])\n",
    "        self.avg_metabolism_seq.append(avg_metabolism)\n",
    "        \n",
    "        # add an agent\n",
    "        add_agents = self.params.get('add_agents', False)\n",
    "        if add_agents:\n",
    "            self.add_agent()\n",
    "        \n",
    "        return len(self.agents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# I'll start with the original model, which does not add agents, and `n=300`\n",
    "\n",
    "RandomSeed(17)\n",
    "\n",
    "env = EvoSugarscape(50, num_agents=300)\n",
    "for i in range(300):\n",
    "    env.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# As we saw before, the number of agents drops to the \"carrying capacity\"\n",
    "\n",
    "thinkplot.plot(env.agent_count_seq)\n",
    "thinkplot.config(ylabel='Agent count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Agents with high metabolism die off quickly, so the average drops quickly.\n",
    "\n",
    "thinkplot.plot(env.avg_metabolism_seq)\n",
    "thinkplot.config(ylabel='Average metabolism')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Agents with poor vision also die quickly, so the average vision increases.\n",
    "\n",
    "thinkplot.plot(env.avg_vision_seq)\n",
    "thinkplot.config(ylabel='Average vision')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Now let's run it again with `add_agents=True`:\n",
    "\n",
    "RandomSeed(17)\n",
    "\n",
    "env = EvoSugarscape(50, num_agents=300, add_agents=True)\n",
    "for i in range(300):\n",
    "    env.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Now the initial population dies off, but then rebounds.\n",
    "\n",
    "# Part of the reason for the increase is that the fitness \n",
    "# of the agents increases, which increases the carrying capacity.\n",
    "\n",
    "# But it is also posssible that the initial die-off overshoots \n",
    "# the actual carrying capacity.\n",
    "\n",
    "thinkplot.plot(env.agent_count_seq)\n",
    "thinkplot.config(ylabel='Agent count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# As we saw before, metabolism drops quickly during the initial die-off, \n",
    "# but now it continues to fall as new agents are born and the least fit die off.\n",
    "\n",
    "thinkplot.plot(env.avg_metabolism_seq)\n",
    "thinkplot.config(ylabel='Average metabolism')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# Average vision increases quickly and then levels off.\n",
    "\n",
    "thinkplot.plot(env.avg_vision_seq)\n",
    "thinkplot.config(ylabel='Average vision')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution\n",
    "\n",
    "# In some runs, the average vision increases at first, \n",
    "# and then falls gradually, which is surprising.\n",
    "\n",
    "# I conjecture that the intial die-off creates unnatural \n",
    "# pressure that kills agents that would have survived in steady state.\n",
    "\n",
    "# After the initial die-off, these agents are able to make \n",
    "# a comeback, which brings the average vision level down.\n",
    "\n",
    "# Also, as the average metabolism continues to fall, \n",
    "# maybe agents can survive despite lower vision."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
